// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <fmt/format.h>
#include <melon/cpp_attributes.h>
#include <limits>
#include <sstream>

#include <pollux/common/base/succinct_printer.h>

namespace kumo::pollux {
    struct RuntimeCounter {
        enum class Unit { kNone, kNanos, kBytes };

        int64_t value;
        Unit unit{Unit::kNone};

        explicit RuntimeCounter(int64_t _value, Unit _unit = Unit::kNone)
            : value(_value), unit(_unit) {
        }
    };

    struct RuntimeMetric {
        // Sum, min, max have the same unit, count has kNone.
        RuntimeCounter::Unit unit;
        int64_t sum{0};
        int64_t count{0};
        int64_t min{std::numeric_limits<int64_t>::max()};
        int64_t max{std::numeric_limits<int64_t>::min()};

        explicit RuntimeMetric(
            RuntimeCounter::Unit _unit = RuntimeCounter::Unit::kNone)
            : unit(_unit) {
        }

        explicit RuntimeMetric(
            int64_t value,
            RuntimeCounter::Unit _unit = RuntimeCounter::Unit::kNone)
            : unit(_unit), sum{value}, count{1}, min{value}, max{value} {
        }

        void addValue(int64_t value);

        /// Aggregate sets 'min' and 'max' to 'sum', also sets 'count' to 1 if
        /// positive.
        void aggregate();

        void printMetric(std::stringstream &stream) const;

        void merge(const RuntimeMetric &other);

        std::string toString() const;
    };

    /// Simple interface to implement writing of runtime stats to Pollux Operator
    /// stats.
    /// Inherit a concrete class from this to implement your writing.
    class BaseRuntimeStatWriter {
    public:
        virtual ~BaseRuntimeStatWriter() = default;

        virtual void addRuntimeStat(
            const std::string & /* name */,
            const RuntimeCounter & /* value */) {
        }
    };

    /// Setting a concrete runtime stats writer on the thread will ensure that any
    /// code can add runtime counters to the current Operator running on that
    /// thread.
    /// NOTE: This is only used by the Pollux Driver at the moment, which ensures the
    /// active Operator is being used by the writer.
    void setThreadLocalRunTimeStatWriter(
        BaseRuntimeStatWriter * MELON_NULLABLE writer);

    /// Retrives the current runtime stats writer.
    BaseRuntimeStatWriter * MELON_NULLABLE getThreadLocalRunTimeStatWriter();

    /// Writes runtime counter to the current Operator running on that thread.
    void addThreadLocalRuntimeStat(
        const std::string &name,
        const RuntimeCounter &value);

    /// Scope guard to conveniently set and revert back the current stat writer.
    class RuntimeStatWriterScopeGuard {
    public:
        explicit RuntimeStatWriterScopeGuard(
            BaseRuntimeStatWriter * MELON_NULLABLE writer)
            : prevWriter_(getThreadLocalRunTimeStatWriter()) {
            setThreadLocalRunTimeStatWriter(writer);
        }

        ~RuntimeStatWriterScopeGuard() {
            setThreadLocalRunTimeStatWriter(prevWriter_);
        }

    private:
        BaseRuntimeStatWriter *const MELON_NULLABLE prevWriter_;
    };
} // namespace kumo::pollux
template<>
struct fmt::formatter<kumo::pollux::RuntimeCounter::Unit> : formatter<int> {
    auto format(kumo::pollux::RuntimeCounter::Unit s, format_context &ctx)
    const {
        return formatter<int>::format(static_cast<int>(s), ctx);
    }
};
